import pandas as pd
import numpy as np
from scipy.stats import norm


def get_resp_cat(num_ques, num_correct):
    if num_ques == 1:
        if num_correct == 0:
            return 1
        elif num_correct == 1:
            return 2
        else:
            raise Exception('总题目数量1，答对题目数量异常')

    elif num_ques == 2:
        if num_correct == 0:
            return 3
        elif num_correct == 1:
            return 4
        elif num_correct == 2:
            return 5
        else:
            raise Exception('总题目数量2，答对题目数量异常')

    elif num_ques == 3:
        if num_correct == 0:
            return 6
        elif num_correct == 1:
            return 7
        elif num_correct == 2:
            return 8
        elif num_correct == 3:
            return 9
        else:
            raise Exception('总题目数量3，大堆题目数量异常')

    elif num_ques == 4:
        if num_correct == 0:
            return 10
        elif num_correct == 1:
            return 11
        elif num_correct == 2:
            return 12
        elif num_correct == 3:
            return 13
        elif num_correct == 4:
            return 14
        else:
            raise Exception('总题目数量4，大堆题目数量异常')

    else:
        raise Exception('总题目数量不是1-4')


def get_item_info(history_at_time_t, max_item_length, max_question_for_item):
    unique_item_history = history_at_time_t.drop_duplicates(subset='item_id', keep='last')
    unique_item_history_desc = unique_item_history.sort_index(ascending=False)

    if len(unique_item_history_desc) > max_item_length:
        items_to_use = unique_item_history_desc[:max_item_length]
    else:
        items_to_use = unique_item_history_desc

    item_resp_cat = np.zeros(len(items_to_use))
    for i, item in enumerate(items_to_use.item_id):
        question_history = history_at_time_t[history_at_time_t.item_id == item]
        unique_question_history = question_history.drop_duplicates(subset='question_id', keep='last')
        unique_question_history_desc = unique_question_history.sort_index(ascending=False)

        if len(unique_question_history_desc) > max_question_for_item:
            questions_to_use = unique_question_history_desc[:max_question_for_item]
        else:
            questions_to_use = unique_question_history_desc

        num_questions = len(questions_to_use)
        num_correctness = questions_to_use.correctness.sum()

        item_resp_cat[i] = get_resp_cat(num_questions, num_correctness)

    item_info = items_to_use.copy()
    item_info.insert(loc=0, column='resp_cat', value=item_resp_cat)

    return item_info


def combine_item_info_param(item_info, item_param, resp_cat_param):
    item_info_with_resp_cat_param = pd.merge(item_info, resp_cat_param,
                                             how='inner', on=['item_id', 'item_textid'])


    num_items = len(item_info)
    resp_cat_multiplier = np.zeros(num_items)

    for i, resp_cat in enumerate(item_info_with_resp_cat_param.resp_cat.values):
        index_for_resp_cat = int(resp_cat) + 5
        resp_cat_multiplier[i] = item_info_with_resp_cat_param.iloc[i, index_for_resp_cat]

    item_info_temp = item_info.copy()
    item_info_temp.insert(loc=0, column='resp_cat_mul', value=resp_cat_multiplier)

    item_info_full = pd.merge(item_info_temp, item_param, how='inner', on=['item_id', 'item_textid'])

    return item_info_full


def calculate_pks(item_info, used_item_param):
    num_cur_items = item_info.shape[0]
    num_used_items = used_item_param.shape[0]

    cur_item_param_slice = slice(8, None)
    used_item_param_slice = slice(3, None)
    item_resp_cat_mul = item_info.resp_cat_mul.values.reshape(num_cur_items, 1)

    cur_item_param = item_info.iloc[:, cur_item_param_slice]
    used_item_center_value = used_item_param.centerValue.values.reshape(num_used_items, 1)
    used_item_param = used_item_param.iloc[:, used_item_param_slice]

    student_shifts = (1 / np.sqrt(num_cur_items)) * np.dot(item_resp_cat_mul.T, cur_item_param)
    item_shifts = np.dot(student_shifts, used_item_param.T).T

    z = item_shifts + used_item_center_value
    return norm.cdf(z)


if __name__ == '__main__':

    MAX_ITEM_HISTORY = 24
    MAX_QUESTION_FOR_ITEM = 4

    student_history = pd.read_csv('./studentHistory.csv')
    question_to_item_and_pks = pd.read_csv('./questionsToItemsMapAndPKS.csv')
    question_to_item_and_pks = question_to_item_and_pks.dropna()
    item_param = pd.read_csv('./itemParametersMap.csv')
    resp_cat_param = pd.read_csv('./itemRespCatMultipliersMap.csv')
    num_question_answered = student_history.shape[0]
    pks_for_questions_before_answered = np.zeros(num_question_answered)
    history_with_item_and_pks_info = pd.merge(student_history, question_to_item_and_pks,
                                              how='inner', on='question_id')
    result = pd.DataFrame(columns=['question_id', 'correctness', 'item_id', 'item_textid', 'PKS', 'predicted_pks'])
    for i in range(1, num_question_answered - 1):
        current_history = history_with_item_and_pks_info[0:i]
        current_item = history_with_item_and_pks_info.iloc[i]
        item_info = get_item_info(current_history, MAX_ITEM_HISTORY, MAX_QUESTION_FOR_ITEM)
        item_info_full = combine_item_info_param(item_info, item_param, resp_cat_param)
        pks = calculate_pks(item_info_full, item_param[item_param.item_id == current_item.item_id])
        pks = pks.reshape(-1)[0]

        pks_to_append = pd.Series([pks], index=['predicted_pks'])
        current_item = current_item.append(pks_to_append)
        result = result.append(current_item, ignore_index=True)

        result.to_csv('./student_history_python.csv')
        question_to_item_and_pks.head()
        history_with_item_and_pks_info.head(20)
